import os

import torch
import torch.nn as nn
from PIL import Image
from torchvision import models, transforms


# 加载类别信息
def load_classes(train_dir):
    classes = sorted(os.listdir(train_dir))
    return classes


# 数据预处理（与训练时一致）
transform = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)


# 预测函数
def predict_images(model, image_paths, classes, device, threshold=0.5):
    model.eval()
    results = []

    for img_path in image_paths:
        try:
            # 加载并预处理图像
            img = Image.open(img_path).convert("RGB")
            img_tensor = transform(img).unsqueeze(0).to(device)

            # 预测
            with torch.no_grad():
                outputs = model(img_tensor)
                probs = torch.nn.functional.softmax(outputs, dim=1)
                max_prob, pred_idx = torch.max(probs, 1)

                # 判断是否为未知类别
                if max_prob.item() < threshold:
                    results.append(
                        {
                            "image": img_path,
                            "prediction": "unknown",
                            "probability": max_prob.item(),
                        }
                    )
                else:
                    results.append(
                        {
                            "image": img_path,
                            "prediction": classes[pred_idx.item()],
                            "probability": max_prob.item(),
                        }
                    )
        except Exception as e:
            print(f"Error processing {img_path}: {str(e)}")
            results.append({"image": img_path, "prediction": "error", "error": str(e)})

    return results


def main():
    # 配置
    model_path = "best_resnet50_model.pth"
    train_dir = "data/TrainingSet"  # 训练数据目录
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 加载类别
    classes = load_classes(train_dir)

    # 加载模型结构
    model = models.resnet50(weights=None)
    model.fc = nn.Linear(2048, len(classes))

    # 加载模型参数
    model.load_state_dict(torch.load(model_path, weights_only=False))
    model = model.to(device)
    model.eval()

    # 获取要预测的图片路径
    input_path = input("请输入要预测的图片路径或目录：")

    # 如果是目录，获取所有图片文件
    if os.path.isdir(input_path):
        image_paths = []
        for root, _, files in os.walk(input_path):
            for file in files:
                if file.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".gif")):
                    image_paths.append(os.path.join(root, file))
    else:
        image_paths = [input_path]

    if not image_paths:
        print("未找到任何图片文件")
        return

    # 进行预测
    results = predict_images(model, image_paths, classes, device)

    # 输出结果
    print("\n预测结果：")
    for result in results:
        if "error" in result:
            print(f"图片: {result['image']}, 错误: {result['error']}")
        else:
            print(
                f"图片: {result['image']}, 预测类别: {result['prediction']}, 置信度: {result['probability']:.4f}"
            )
            


if __name__ == "__main__":
    main()
